from __future__ import annotations
from dataclasses import dataclass
from .utils.logger import get_logger
from .dto import mqtt
from .devices import CONNECTIONS, ClientDevice

from amqtt.broker import Broker, BrokerContext
from amqtt.plugins.base import BasePlugin, BaseTopicPlugin
from amqtt.session import Session, IncomingApplicationMessage
import asyncio
import logging
import traceback

"""
from amqtt import events

包含的event定义如下:
class MQTTEvents(Events):
    PACKET_SENT = "mqtt_packet_sent"
    PACKET_RECEIVED = "mqtt_packet_received"

class BrokerEvents(Events):
    PRE_START = "broker_pre_start"
    POST_START = "broker_post_start"
    PRE_SHUTDOWN = "broker_pre_shutdown"
    POST_SHUTDOWN = "broker_post_shutdown"
    CLIENT_CONNECTED = "broker_client_connected"
    CLIENT_DISCONNECTED = "broker_client_disconnected"
    CLIENT_SUBSCRIBED = "broker_client_subscribed"
    CLIENT_UNSUBSCRIBED = "broker_client_unsubscribed"
    RETAINED_MESSAGE = "broker_retained_message"
    MESSAGE_RECEIVED = "broker_message_received"
    MESSAGE_BROADCAST = "broker_message_broadcast"

from amqtt import broker # 在amqtt/broker.py下搜索`fire_event`获取具体Events用法(参数定义)
"""

logger = get_logger(__name__, level=logging.DEBUG)
# %% Mqtt Broker配置
config = {
    "listeners": {
        "default": {
            "type": "tcp",
            "bind": "0.0.0.0:1883",
        },
        # "ws-mqtt": {
        #     "bind": "127.0.0.1:8080",
        #     "type": "ws",
        #     "max_connections": 10,
        # },
    },
    "plugins": {
        # 允许匿名连接
        "amqtt.plugins.authentication.AnonymousAuthPlugin": {"allow_anonymous": True},
        # "amqtt.plugins.sys.broker.BrokerSysPlugin": {
        #     "sys_interval": 10
        # },  # 定时broadcast
        # "amqtt.plugins.topic_checking.TopicTabooPlugin": {},
        "src.broker.CustomPlugin": {
            "display_port": True,
            "connect_resp": "hello",
        },
    },
}
# %% 自定义插件(处理Broker事件)
class CustomPlugin(BasePlugin):
    # defined slots type
    config: CustomPlugin.Config
    context: BrokerContext

    @dataclass
    class Config:
        display_port: bool = False
        connect_resp: bytes | str = ""
        count: int = 0  # 定时广播计数

        def __post_init__(self):
            if isinstance(self.connect_resp, str):
                self.connect_resp = self.connect_resp.encode()

    async def send_to_session(self, session: Session, topic: str, data: bytes):
        handler = self.context._broker_instance._get_handler(session)
        if handler:
            await handler.mqtt_publish(topic, data, qos=0, retain=False)
            logger.debug(
                f">> Message sent to session {session.client_id} on topic '{topic}, {data}'"
            )
        else:
            logger.warning(
                f">> Handler for session {session.client_id} not found, cannot send message."
            )

    async def on_broker_post_start(self) -> None:  # 启动定时broadcast任务
        broadcast_msg = mqtt.CustomMessage(
            payload={"message": f"hello with count {self.config.count}"},
        ).to_bytes()
        self.config.count += 1
        tasks = [
            self.send_to_session(
                device.client_session, device.subscribe_topic, broadcast_msg
            )
            for device in CONNECTIONS.values()
        ] # 该发送为强制性的, 无需对方订阅
        tasks.append(
            self.context.broadcast_message("broker/broadcast", broadcast_msg, 0)
        )  # 需要对方subscribe broadcast
        if tasks:
            await asyncio.gather(*tasks)
            logger.debug("Broadcast 'ping' message to all connected clients.")
        self.context.loop.call_later(
            10, lambda: self.context.loop.create_task(self.on_broker_post_start())
        )  # 定时

    async def on_broker_client_connected(
        self, *, client_id: str, client_session: Session
    ) -> None:
        cfg = self.config
        display_port_str = (
            f"on port '{client_session.remote_port}'" if cfg.display_port else ""
        )
        client_session.keep_alive = 60
        logger.info(
            f"<< client '{client_id}' connected from '{client_session.remote_address}' {display_port_str}"
        )
        CONNECTIONS[client_id] = ClientDevice(client_id, client_session)

    async def on_broker_message_received(
        self, *, client_id: str, message: IncomingApplicationMessage
    ):
        """
        message slots
        ("data",
        "packet_id",
        "puback_packet",
        "pubcomp_packet",
        "publish_packet",
        "pubrec_packet",
        "pubrel_packet",
        "qos",
        "retain",
        "topic")
        """
        logger.info(
            f"<< Message received from client '{client_id}':{message.topic}:{message.data}"
        )
        # TODO 校验后将消息转发给订阅该topic的设备
        try:
            struct = mqtt.MessageDecoder.from_bytes(message.data)
            if isinstance(struct, mqtt.HelloMessage):
                device = CONNECTIONS.get(client_id)
                if device is None:
                    logger.warning(f"Client '{client_id}' not found in connections.")
                    return
                resp = device.from_hello(struct).hello_back().to_bytes()
                await self.send_to_session(
                    device.client_session, device.subscribe_topic, resp
                )
        except Exception:
            logger.error(
                f"Failed to decode message from client '{client_id}': {traceback.format_exc()}"
            )

    async def on_broker_client_disconnected(
        self, *, client_id: str, client_session: Session
    ):
        device = CONNECTIONS.pop(client_id, None)
        if device:
            del device  # 清理设备引用
            logger.info(f">> Client '{client_id}' disconnected.")
        else:
            logger.warning(f">> Client '{client_id}' not found in connections.")
        del client_session  # 清理会话引用

# %% 启动Broker
async def start():
    broker = Broker(config)
    try:
        await broker.start()
        await asyncio.Event().wait()
    except asyncio.CancelledError:
        await broker.shutdown()
